#!/usr/bin/python3 python3
# -*- coding:utf-8 -*-
# @FileName  :Mrosdep.py
# @Author    :TFly
import os
import subprocess
from colorama import Fore,Style

url = "https://gitee.com/tyx6/rosdistro/raw/master"
url_list = "https://gitee.com/tyx6/mytools/raw/main/ros/20-default.list" 
ros_distros = ['kinetic', 'melodic', 'noetic', 'foxy', 'galactic', 'humble', 'iron', 'rolling']

def find_ros_setup_script():
    #ros_distros = ['melodic', 'noetic']
    for distro in ros_distros:
        setup_script = f"/opt/ros/{distro}/setup.bash"
        if os.path.exists(setup_script):
            return setup_script
    return None

def get_ros_version():
    try:
        ros_setup_script = find_ros_setup_script()
        if not ros_setup_script:
            print("ROS setup script not found.")
            return None

        command = f"source {ros_setup_script} && echo $ROS_DISTRO"
        ros_version = subprocess.check_output(['bash', '-c', command], universal_newlines=True).strip()
        return ros_version
    except subprocess.CalledProcessError as e:
        print(f"An error occurred while getting ROS version: {e}")
        return None

def rep(data):
    flag_list = data.find("https://raw.githubusercontent.com/ros/rosdistro/master/rosdep/sources.list.d/20-default.list")
    if flag_list != -1:
        data = data.replace("https://raw.githubusercontent.com/ros/rosdistro/master/rosdep/sources.list.d/20-default.list",url_list)

    flag = data.find("https://raw.githubusercontent.com/ros/rosdistro/master")
    flag1 = data.find("FUERTE_GBPDISTRO_URL = \'https://raw.githubusercontent.com/ros/rosdistro")
    flag2 = data.find("master/releases/fuerte.yaml")
    if flag != -1:
        data = data.replace("https://raw.githubusercontent.com/ros/rosdistro/master", url)
    if flag1 != -1:
        data = data.replace("https://raw.githubusercontent.com/ros/rosdistro",url)
    if flag2 != -1:
        data = data.replace("master/releases/fuerte.yaml","releases/fuerte.yaml")
    
    return data

def re_files(fileName):
    try:
        con = ""
        with open(fileName,'r') as fp:
            con = fp.readlines()
            for i in range(0,len(con)):
                con[i] = rep(con[i])

        with open(fileName,'w+') as fp:
            for str1 in con:
                fp.write(str1)
    except FileNotFoundError:
        print(fileName+'not found')
    else:
        print(fileName+' is ok')


if __name__ == '__main__':
    print(Fore.GREEN + '欢迎使用脚本，本脚本适用于ROS1、ROS2.\n'+
            'Test: ROS1 Melodic、ROS2 Noetic、ROS2 Humble' + Style.RESET_ALL)
    version = get_ros_version()

    if version not in ros_distros :
        print(Fore.RED + "不支持此版本、或版本获取失败！"+ Style.RESET_ALL)
    else:
        print(Fore.GREEN + "The ROS version is: " + version + Style.RESET_ALL)
        fileName = []

        if version == 'melodic' :
            fileName = ['/usr/lib/python2.7/dist-packages/rosdep2/sources_list.py',
                    '/usr/lib/python2.7/dist-packages/rosdistro/__init__.py',
                    '/usr/lib/python2.7/dist-packages/rosdep2/gbpdistro_support.py',
                    '/usr/lib/python2.7/dist-packages/rosdep2/rep3.py']
        else:
            fileName = ['/usr/lib/python3/dist-packages/rosdep2/sources_list.py',
                    '/usr/lib/python3/dist-packages/rosdistro/__init__.py',
                    '/usr/lib/python3/dist-packages/rosdep2/gbpdistro_support.py',
                    '/usr/lib/python3/dist-packages/rosdep2/rep3.py']

            

        for name in fileName:
            re_files(name)

        print("修改成功！请执行下面命令：\n"+
                "sudo rosdep init\n"+
                "rosdep update")
